﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using TorchSharp;
using TorchSharp.Modules;
using static TorchSharp.torch;

namespace NgpSharp
{
    public class HashEmbedder : nn.Module<Tensor, int, Tuple<Tensor, Tensor>>
    {
        public int OutDim { get; private set; }

        protected int n_levels;
        protected int n_features_per_level;
        protected int log2_hashmap_size;
        protected Tensor base_resolution;
        protected Tensor finest_resolution;
        protected Tensor b;
        protected ModuleList<Embedding> embeddings;
        public HashEmbedder(string name, int n_levels = 16, 
            int n_features_per_level = 2, int log2_hashmap_size = 19,
            int base_resolution = 16, int finest_resolution = 512) 
            : base(name)
        {
            this.n_levels = n_levels;
            this.n_features_per_level = n_features_per_level;
            this.log2_hashmap_size = log2_hashmap_size;
            this.base_resolution = torch.tensor(base_resolution);
            this.finest_resolution = torch.tensor(finest_resolution);
            OutDim = n_levels * n_features_per_level;
            b = torch.exp((torch.log(finest_resolution) - torch.log(base_resolution)) / (n_levels - 1));
            List<Embedding> embeddings_list = new List<Embedding>();
            for (int i = 0; i < n_levels; i++) 
            {
                var e = nn.Embedding((long)Math.Pow(2, log2_hashmap_size), n_features_per_level);
                embeddings_list.Add(e);
            }
            embeddings = nn.ModuleList<Embedding>(embeddings_list.ToArray());
            for(int i = 0; i < n_levels; i++)
            {
                //nn.init.constant_(embeddings[i].weight, 0.1 * i);
                nn.init.uniform_(embeddings[i].weight, low: -0.0001, high: 0.0001);
            }
            RegisterComponents();
        }

        public override Tuple<Tensor, Tensor> forward(Tensor x, int bound)
        {
            using var c = torch.NewDisposeScope();
            List<Tensor> x_embedded_all = [];
            Tensor keep_mask = false;

            Tensor box_min = torch.from_array(new float[] { -bound, -bound, -bound }, device: x.device);
            Tensor box_max = torch.from_array(new float[] { bound, bound, bound }, device: x.device);
            if (!torch.all(x <= box_max).ToBoolean() || !torch.all(x >= box_min).ToBoolean())
            {
                x = torch.clamp(x, min: box_min, max: box_max);
            }

            for (int i = 0; i < n_levels; i++)
            {
                Tensor resolution = torch.floor(base_resolution * torch.pow(b, i));
                (Tensor voxel_min_vertex, Tensor voxel_max_vertex,
                    Tensor hashed_voxel_indices, keep_mask) = GetVoxelVertices(x,
                    box_min, box_max, resolution, log2_hashmap_size);
                Tensor voxel_embedds = embeddings[i].call(hashed_voxel_indices);

                Tensor x_embedded = TrilinearInterp(x, voxel_min_vertex, voxel_max_vertex, voxel_embedds);
                x_embedded_all.Add(x_embedded);

                keep_mask = keep_mask.sum(dim: -1) == keep_mask.shape[^1];
            }
            return Tuple.Create(torch.cat(x_embedded_all, dim: -1).MoveToOuterDisposeScope(), 
                keep_mask.MoveToOuterDisposeScope());
        }

        private Tuple<Tensor, Tensor, Tensor, Tensor> GetVoxelVertices(Tensor xyz, Tensor box_min, Tensor box_max, Tensor resolution, double log2_hashmap_size)
        {
            using var c = torch.NewDisposeScope();

            Tensor keep_mask = xyz == torch.max(torch.min(xyz, box_max), box_min);

            Tensor grid_size = (box_max - box_min) / resolution;

            Tensor bottom_left_idx = torch.floor((xyz - box_min) / grid_size).to(ScalarType.Int32);
            Tensor voxel_min_vertex = bottom_left_idx * grid_size + box_min;
            Tensor voxel_max_vertex = voxel_min_vertex + torch.tensor(new float[] { 1.0f, 1.0f, 1.0f }, device: xyz.device) * grid_size;
            Tensor voxel_indices = bottom_left_idx.unsqueeze(1) + torch.from_array(OFFSET, device: xyz.device);

            Tensor hashed_voxel_indices = Hash(voxel_indices, (int)log2_hashmap_size);

            return Tuple.Create(voxel_min_vertex.MoveToOuterDisposeScope(),
                voxel_max_vertex.MoveToOuterDisposeScope(),
                hashed_voxel_indices.MoveToOuterDisposeScope(),
                keep_mask.MoveToOuterDisposeScope());
        }

        private static readonly long[] PRIMES = [1, 2654435761, 805459861, 3674653429, 2097192037, 1434869437, 2165219737];
        private static readonly int[,] OFFSET = GetOffset();

        private Tensor Hash(Tensor coords, int log2_hashmap_size)
        {
            using var c = torch.NewDisposeScope();
            Tensor xor_result = torch.zeros_like(coords)[TensorIndex.Ellipsis, 0];
            for(int i = 0; i < coords.shape[^1]; i++)
            {
                xor_result ^= coords[TensorIndex.Ellipsis, i] * PRIMES[i];
            }
            Tensor r = torch.tensor((1 << log2_hashmap_size) - 1).to(xor_result.device) & xor_result;
            return r.MoveToOuterDisposeScope();
        }

        private Tensor TrilinearInterp(Tensor x, Tensor voxel_min_vertex, Tensor voxel_max_vertex, Tensor voxel_embedds)
        {
            using var context = torch.NewDisposeScope();
            Tensor weights = (x - voxel_min_vertex) / (voxel_max_vertex - voxel_min_vertex);

            Tensor c00 = voxel_embedds[TensorIndex.Colon, 0] * (1 - weights[TensorIndex.Colon, 0][TensorIndex.Colon, TensorIndex.None]) + voxel_embedds[TensorIndex.Colon, 4] * weights[TensorIndex.Colon, 0][TensorIndex.Colon, TensorIndex.None];
            Tensor c01 = voxel_embedds[TensorIndex.Colon, 1] * (1 - weights[TensorIndex.Colon, 0][TensorIndex.Colon, TensorIndex.None]) + voxel_embedds[TensorIndex.Colon, 5] * weights[TensorIndex.Colon, 0][TensorIndex.Colon, TensorIndex.None];
            Tensor c10 = voxel_embedds[TensorIndex.Colon, 2] * (1 - weights[TensorIndex.Colon, 0][TensorIndex.Colon, TensorIndex.None]) + voxel_embedds[TensorIndex.Colon, 6] * weights[TensorIndex.Colon, 0][TensorIndex.Colon, TensorIndex.None];
            Tensor c11 = voxel_embedds[TensorIndex.Colon, 3] * (1 - weights[TensorIndex.Colon, 0][TensorIndex.Colon, TensorIndex.None]) + voxel_embedds[TensorIndex.Colon, 7] * weights[TensorIndex.Colon, 0][TensorIndex.Colon, TensorIndex.None];

            Tensor c0 = c00 * (1 - weights[TensorIndex.Colon, 1][TensorIndex.Colon, TensorIndex.None]) + c10 * weights[TensorIndex.Colon, 1][TensorIndex.Colon, TensorIndex.None];
            Tensor c1 = c01 * (1 - weights[TensorIndex.Colon, 1][TensorIndex.Colon, TensorIndex.None]) + c11 * weights[TensorIndex.Colon, 1][TensorIndex.Colon, TensorIndex.None];

            Tensor c = c0 * (1 - weights[TensorIndex.Colon, 2][TensorIndex.Colon, TensorIndex.None]) + c1 * weights[TensorIndex.Colon, 2][TensorIndex.Colon, TensorIndex.None];
            return c.MoveToOuterDisposeScope();
        }

        private static int[,] GetOffset()
        {
            int[,] offsets = new int[8, 3];
            int index = 0;
            for (int i = 0; i <= 1; i++)
            {
                for (int j = 0; j <= 1; j++)
                {
                    for (int k = 0; k <= 1; k++)
                    {
                        offsets[index, 0] = i;
                        offsets[index, 1] = j;
                        offsets[index, 2] = k;
                        index++;
                    }
                }
            }
            return offsets;
        }
    }
}
